"""
CN-CLIP 预训练模型加载器
"""

import torch
import cn_clip.clip as clip
from cn_clip.clip import load_from_name, available_models
import logging

logger = logging.getLogger(__name__)


class CNClipModelLoader:
    """CN-CLIP预训练模型加载和管理"""
    
    def __init__(self, device="cuda" if torch.cuda.is_available() else "cpu"):
        self.device = device
        self.model = None
        self.preprocess = None
        self.model_name = None
        
    def load_model(self, model_name="ViT-B-16", download_root="./pretrained_weights"):
        """
        加载CN-CLIP预训练模型
        
        Args:
            model_name: 模型名称，如 'ViT-B-16', 'ViT-L-14' 等
            download_root: 模型下载存储路径
        """
        try:
            logger.info(f"Loading CN-CLIP model: {model_name}")
            self.model, self.preprocess = load_from_name(
                model_name, 
                device=self.device, 
                download_root=download_root
            )
            self.model_name = model_name
            self.model.eval()  # 设置为评估模式
            logger.info(f"Model {model_name} loaded successfully on {self.device}")
            return True
        except Exception as e:
            logger.error(f"Failed to load model {model_name}: {e}")
            return False
    
    def get_available_models(self):
        """获取所有可用的CN-CLIP模型列表"""
        return available_models()
    
    def encode_image(self, images):
        """编码图像特征"""
        if self.model is None:
            raise RuntimeError("Model not loaded. Call load_model() first.")
        
        with torch.no_grad():
            image_features = self.model.encode_image(images)
            # L2归一化
            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
        return image_features
    
    def encode_text(self, text_tokens):
        """编码文本特征"""
        if self.model is None:
            raise RuntimeError("Model not loaded. Call load_model() first.")
            
        with torch.no_grad():
            text_features = self.model.encode_text(text_tokens)
            # L2归一化
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
        return text_features
    
    def compute_similarity(self, image_features, text_features):
        """计算图像-文本相似度矩阵"""
        # 计算余弦相似度
        similarity = image_features @ text_features.T
        return similarity
    
    def get_model_info(self):
        """获取模型信息"""
        if self.model is None:
            return "No model loaded"
        
        total_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        
        return {
            "model_name": self.model_name,
            "device": self.device,
            "total_parameters": total_params,
            "trainable_parameters": trainable_params,
            "is_training": self.model.training
        }


if __name__ == "__main__":
    # 测试模型加载
    loader = CNClipModelLoader()
    print("Available models:", loader.get_available_models())
    
    if loader.load_model("ViT-B-16"):
        print("Model info:", loader.get_model_info())
    else:
        print("Failed to load model")